import numpy as np
import scipy.special as sc
from chaco.api import *
from chaco.tools.api import DragZoom, PanTool, ZoomTool
from enable.api import ComponentEditor
from traits.api import *
from traitsui.api import *

depthp = np.array([0.0,
                  200.0,
                  250.0,
                  400.0,
                  600.0,
                  800.0,
                  1000.0,
                  1200.0,
                  1400.0,
                  1600.0,
                  1800.0,
                  2000.0,
                  2200.0,
                  2400.0,
                  2600.0,
                  2800.0,
                  3000.0,
                  3200.0,
                  3400.0,
                  3600.0,
                  3800.0,
                  4000.0,
                  4200.0,
                  4400.0,
                  4600.0,
                  4800.0,
                  5000.0, ])
speedp = np.array([
    1548.52,
    1530.29,
    1526.69,
    1517.78,
    1509.49,
    1504.30,
    1501.38,
    1500.14,
    1500.12,
    1501.02,
    1502.57,
    1504.62,
    1507.02,
    1509.69,
    1512.55,
    1515.56,
    1518.67,
    1521.85,
    1525.10,
    1528.38,
    1531.70,
    1535.04,
    1538.39,
    1541.76,
    1545.14,
    1548.52,
    1551.91, ])
def speed(depth):
    return np.interp(depth,depthp,speedp)
def gradient(depth):
    return speed(depth+1)-speed(depth)

def ray(depth=500,angle=0,N=10,step=100):
    pt=np.zeros((N+1,2))
    pt[0,1]=depth
    n=np.array([np.cos(angle*np.pi/180),np.sin(angle*np.pi/180)])
    for k in range(1,N+1):
        pt[k]=pt[k-1]+step*n
        n+=np.array([n[0]*n[1],n[1]**2-1])*gradient(pt[k-1,1])/speed(pt[k-1,1])
        n/=np.sum(n*n)**.5
        if pt[k,1]<0 or pt[k,1]>5000:
            n[1]=-n[1]
    return pt

# import matplotlib.pyplot as plt
# fig=plt.figure()
# ax=fig.add_subplot()
# pt0=ray(500,0,5000,500)
# ax.plot(pt0[:,0],pt0[:,1],label='0 degree')
# pt10=ray(500,0.5,5000,500)
# ax.plot(pt10[:,0],pt10[:,1],label='500m,10 degree')
# pt=ray(1000,0.5,5000,500)
# ax.plot(pt[:,0],pt[:,1],label='1000m,10 degree')
# pt=ray(1000,2,5000,500)
# ax.plot(pt[:,0],pt[:,1],label='1000m,10 degree',marker='o')
# ax.set_ylim([0,8000])
# ax.invert_yaxis()
#
# plt.show()
#
#
#
# import sys
# sys.exit()

class BasicRelation(HasTraits):
    plot = Instance(Plot)

    def _plot_default(self):
        plot = Plot(ArrayPlotData())
        self.clear_plot(plot)
        return plot

    fun = Enum('声速剖面', )

    def _fun_changed(self):
        plot = self.plot
        self.clear_plot(plot)
        match self.fun:
            case '声速剖面':
                plot.data.update_data(
                    depth=depthp,
                    speed=speedp,
                )
                plot.x_axis.title = "speed [m/s]"
                plot.y_axis.title = "depth [m]"
                plot.orientation='v'
                plot.plot(('depth', 'speed'), name='Speed Profile', color='auto')
        legend = Legend()
        legend.plots = plot.plots
        plot.overlays.append(legend)
        plot.request_redraw()

    def traits_init(self):
        self._fun_changed()

    def clear_plot(self, plot):
        plots = [p for p in plot.plots.keys()]
        for p in plots:
            plot.delplot(p)
        for overlay in plot.overlays:
            plot.overlays.remove(overlay)
        plot.request_redraw()

    traits_view = View(
        Group(
            UCustom(
                "plot",
                editor=ComponentEditor(),
                resizable=True
            ),
            HGroup(
                UItem('fun'),
                spring,
            ),
        ),
        width=1000,
        height=800,
        resizable=True,
    )

class RayTrace(HasTraits):
    plot = Instance(Plot)

    def _plot_default(self):
        plot = Plot(ArrayPlotData())
        self.clear_plot(plot)
        return plot

    fun = Enum('单射线', '多射线')

    depth=Range(0,5000.,1000)
    angle=Range(-5,5.,0)
    number=Range(100,5000,1000)
    step=Range(100,2000,500)

    def _fun_changed(self):
        plot = self.plot
        self.clear_plot(plot)
        match self.fun:
            case '单射线':
                pt=ray(self.depth,self.angle,self.number,self.step)
                plot.data.update_data(
                    x=pt[:,0],
                    y=pt[:,1],
                )
                plot.x_axis.title = "distance [m]"
                plot.y_axis.title = "depth [m]"
                plot.value_range.low_setting = 0
                plot.value_range.high_setting = 5000
                plot.plot(('x', 'y'), name='Ray', color='auto')
            case '多射线':
                pt=ray(self.depth,0,self.number,self.step)
                plot.data.update_data(
                    x0=pt[:,0],
                    y0=pt[:,1],
                )
                pt=ray(self.depth,1,self.number,self.step)
                plot.data.update_data(
                    x1=pt[:,0],
                    y1=pt[:,1],
                )
                pt=ray(self.depth,2,self.number,self.step)
                plot.data.update_data(
                    x2=pt[:,0],
                    y2=pt[:,1],
                )
                pt=ray(self.depth,-1,self.number,self.step)
                plot.data.update_data(
                    x1_=pt[:,0],
                    y1_=pt[:,1],
                )
                pt=ray(self.depth,-2,self.number,self.step)
                plot.data.update_data(
                    x2_=pt[:,0],
                    y2_=pt[:,1],
                )
                plot.x_axis.title = "distance [m]"
                plot.y_axis.title = "depth [m]"
                plot.value_range.low_setting = 0
                plot.value_range.high_setting = 5000
                plot.plot(('x0', 'y0'), name='0 deg', color='auto')
                plot.plot(('x1', 'y1'), name='1 deg', color='auto')
                plot.plot(('x2', 'y2'), name='2 deg', color='auto')
                plot.plot(('x1_', 'y1_'), name='-1 deg', color='auto')
                plot.plot(('x2_', 'y2_'), name='-2 deg', color='auto')
        legend = Legend()
        legend.plots = plot.plots
        plot.overlays.append(legend)
        plot.request_redraw()

    @observe('depth,angle,number,step')
    def update(self,event=None):
        match self.fun:
            case '单射线':
                pt=ray(self.depth,self.angle,self.number,self.step)
                self.plot.data.update_data(
                    x=pt[:,0],
                    y=pt[:,1],
                )
            case '多射线':
                pt=ray(self.depth,0,self.number,self.step)
                self.plot.data.update_data(
                    x0=pt[:,0],
                    y0=pt[:,1],
                )
                pt=ray(self.depth,1,self.number,self.step)
                self.plot.data.update_data(
                    x1=pt[:,0],
                    y1=pt[:,1],
                )
                pt=ray(self.depth,2,self.number,self.step)
                self.plot.data.update_data(
                    x2=pt[:,0],
                    y2=pt[:,1],
                )
                pt=ray(self.depth,-1,self.number,self.step)
                self.plot.data.update_data(
                    x1_=pt[:,0],
                    y1_=pt[:,1],
                )
                pt=ray(self.depth,-2,self.number,self.step)
                self.plot.data.update_data(
                    x2_=pt[:,0],
                    y2_=pt[:,1],
                )
    def traits_init(self):
        self._fun_changed()

    def clear_plot(self, plot):
        plots = [p for p in plot.plots.keys()]
        for p in plots:
            plot.delplot(p)
        for overlay in plot.overlays:
            plot.overlays.remove(overlay)
        plot.request_redraw()

    traits_view = View(
        Group(
            UCustom(
                "plot",
                editor=ComponentEditor(),
                resizable=True
            ),
            HGroup(
                UItem('fun'),
                spring,
            ),
            Item('depth', label='声源深度[m]', editor_args={'mode': 'slider'}),
            Item('angle', label='发射角度[deg]', editor_args={'mode': 'slider'}),
            Item('number', label='步数', editor_args={'mode': 'slider'}),
            Item('step', label='射线步长[m]', editor_args={'mode': 'slider'}),
        ),
        width=1000,
        height=800,
        resizable=True,
    )


class Form(HasTraits):
    ray_trace = Instance(RayTrace, ())
    basic_relation = Instance(BasicRelation, ())
    traits_view = View(
        Tabbed(
            UCustom('ray_trace', label='声射线'),
            UCustom('basic_relation', label='声速剖面'),
        ),
        width=1000,
        height=800,
        resizable=True,
    )


if __name__ == '__main__':
    Form().configure_traits()
